from get_model.model_computation_generator_by_layer import ModelGeneratorByLayer
from get_model.parse_hardware_config import parse_hardware_config
from solver.intra_solver_by_ILP.cluster_env import ClusterEnvironment
from solver.intra_solver_by_ILP.hloLayer import *
from solver.intra_solver_by_ILP.ILP_solver import SolverOption, solve_layer_auto_sharding, \
    solve_layer_auto_sharding_with_resharding, solve_layer_auto_sharding_with_resharding_without_print

MB = 1024 * 1024
hardware_config_path = "/home/lthpc/nvmessd/wangjiaqian/simulator_by_alpa_layer/config/hardware_config.yaml"

_, hardware_config_device_mesh = parse_hardware_config(hardware_config_path)

def assert_close(x, y):
    assert abs(x / y - 1) < 0.001, f"{x} vs. {y}"


def solve_without_all_gather(layerComputation, submesh_shape):
    """ 解决 ILP 问题，不使用 All-Gather 操作 """
    print(submesh_shape)
    sub_device_mesh = np.arange(np.prod(submesh_shape)).reshape(submesh_shape)  # 生成设备网格
    solver_option = SolverOption()
    # solver_option.force_all_gather_cost = 1e8  # 禁止使用 all-gather
    cluster_env = ClusterEnvironment(sub_device_mesh,
                                     [1, 1],
                                     [1, 1],
                                     memory_per_device=8000 * 1024 * MB,
                                     solver_option=solver_option)
    # 求解 ILP 并返回优化目标
    objective = solve_layer_auto_sharding_with_resharding(layerComputation, cluster_env, solver_option)
    return objective

def solve_without_all_gather_without_print(layerComputation, submesh_shape):
    """ 解决 ILP 问题，不使用 All-Gather 操作 """
    print(submesh_shape)
    sub_device_mesh = np.arange(np.prod(submesh_shape)).reshape(submesh_shape)  # 生成设备网格
    solver_option = SolverOption()
    # solver_option.force_all_gather_cost = 1e8  # 禁止使用 all-gather
    cluster_env = ClusterEnvironment(sub_device_mesh,
                                     [1, 1],
                                     [1, 1],
                                     memory_per_device=8000 * 1024 * MB,
                                     solver_option=solver_option)
    # 求解 ILP 并返回优化目标
    objective = solve_layer_auto_sharding_with_resharding_without_print(layerComputation, cluster_env, solver_option)
    return objective


def test_print():
    """ 打印 Layer 计算图 """
    modelGeneratorByLayer = ModelGeneratorByLayer()
    layerComputation = modelGeneratorByLayer.computation_graph(modelGeneratorByLayer.module_length)
    print("===== Hlo Layer Computation =====")
    print(layerComputation, "\n")

def test_stage_print():
    """ 打印 stage 计算图 """
    modelGeneratorByLayer = ModelGeneratorByLayer()
    end = modelGeneratorByLayer.module_length
    layerComputation = modelGeneratorByLayer.computation_stage_graph(0, end)
    print("===== Hlo Layer Computation =====")
    print(layerComputation, "\n")


def test_model_solve():
    """ 运行模型求解 """
    sub_mesh_shape = hardware_config_device_mesh  # 设备网格
    modelGeneratorByLayer = ModelGeneratorByLayer()
    layerComputation = modelGeneratorByLayer.computation_graph(modelGeneratorByLayer.module_length)

    # 计算 ILP 求解目标
    objective = solve_without_all_gather(layerComputation, sub_mesh_shape)

    # print("---------------cluster_env--------------")
    # print(cluster_env.device_mesh)
    print("Objective:", objective)

def test_model_stage_solve():
    """ 调用computation_stage_graph() 方法"""
    # sub_mesh_shape = hardware_config_device_mesh  # 设备网格
    sub_mesh_shape_test = [1, 4]
    modelGeneratorByLayer = ModelGeneratorByLayer()
    end = modelGeneratorByLayer.module_length
    layerComputation = modelGeneratorByLayer.computation_stage_graph(0, end)
    layerComputation_test = modelGeneratorByLayer.computation_graph(3)
    # 计算 ILP 求解目标
    objective = solve_without_all_gather(layerComputation, sub_mesh_shape_test)

    # print("---------------cluster_env--------------")
    # print(cluster_env.device_mesh)
    print("Objective:", objective)

# best solution is :  [((0, 2), (4, 8), 1), ((2, 9), (1, 8), 1), ((9, 16), (1, 8), 1), ((16, 18), (1, 8), 1), ((18, 21), (1, 8), 1)]
def test_model_stage_solve_dp(current_module, next_start_module, submesh_shape):
    modelGeneratorByLayer = ModelGeneratorByLayer()
    layerComputation = modelGeneratorByLayer.computation_stage_graph(current_module, next_start_module)
    # 计算 ILP 求解目标
    objective = solve_without_all_gather_without_print(layerComputation, submesh_shape)

    # print("---------------cluster_env--------------")
    # print(cluster_env.device_mesh)
    print("Objective:", objective)


if __name__ == '__main__':
    # 直接执行测试代码
    # test_print()
    # test_stage_print()
    # test_m odel_solve()
    test_model_stage_solve()
